package org.sean.framework.canary.loadbalancer;

import lombok.extern.slf4j.Slf4j;
import org.sean.framework.canary.loadbalancer.predicate.AbstractServerPredicate;
import org.sean.framework.canary.loadbalancer.predicate.CompositePredicate;
import org.sean.framework.canary.loadbalancer.predicate.LabelPredicate;
import org.sean.framework.canary.loadbalancer.predicate.LocalLabelPredicate;
import org.sean.framework.canary.loadbalancer.predicate.MainPredicate;
import org.sean.framework.canary.loadbalancer.predicate.MasterPredicate;
import org.sean.framework.canary.loadbalancer.predicate.ServerPredicate;
import org.sean.framework.canary.loadbalancer.selector.BestAvailableWeightedSelector;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.loadbalancer.core.DelegatingServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.http.HttpHeaders;
import reactor.core.publisher.Flux;

import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * Label选择器
 */
@Slf4j
public class CanaryServicesSupplier extends DelegatingServiceInstanceListSupplier {

    // 灰度选择器
    private final AbstractServerPredicate predicate;
    // 最优选择器
    private final BestAvailableWeightedSelector bestAvailableWeightedSelector;

    public CanaryServicesSupplier(ServiceInstanceListSupplier delegate) {
        super(delegate);
        this.predicate = createPredicate();
        bestAvailableWeightedSelector = new BestAvailableWeightedSelector();
    }

    protected CompositePredicate createPredicate() {
        // Label过滤器
        LabelPredicate labelPredicate = new LabelPredicate(1);
        LabelPredicate labelPredicate2Level = new LabelPredicate(2);
        // 同版本过滤
        LocalLabelPredicate localLabelPredicate = new LocalLabelPredicate();
        // 自定义默认版本
        MainPredicate mainPredicate = new MainPredicate();
        // 兜底版本
        MasterPredicate masterPredicate = new MasterPredicate();
        // 指定服务
        ServerPredicate serverPredicate = new ServerPredicate();

        return CompositePredicate.withPredicates(labelPredicate)
                .addFallbackPredicate(CompositePredicate.withPredicates(labelPredicate2Level).build())
                .addFallbackPredicate(CompositePredicate.withPredicates(localLabelPredicate).build())
                .addFallbackPredicate(CompositePredicate.withPredicates(serverPredicate).build())
                .addFallbackPredicate(CompositePredicate.withPredicates(mainPredicate).build())
                .addFallbackPredicate(CompositePredicate.withPredicates(masterPredicate).build())
                .build();
    }

    @Override
    public String getServiceId() {
        return delegate.getServiceId();
    }

    @Override
    public Flux<List<ServiceInstance>> get() {
        return delegate.get();
    }

    @SuppressWarnings("rawtypes")
    @Override
    public Flux<List<ServiceInstance>> get(Request request) {
        Object context = request.getContext();
        if ((context instanceof RequestDataContext)) {
            return get().map(serviceInstances -> selectInstance(serviceInstances,
                    ((RequestDataContext) context).getClientRequest().getHeaders()));
        }
        // If the object type is not RequestData, we return all the instances provided by
        // the delegate.
        return get();
    }

    /**
     * 选择服务
     *
     * @param serviceInstances 服务列表
     * @param headers          请求头
     * @return 服务列表
     */
    protected List<ServiceInstance> selectInstance(List<ServiceInstance> serviceInstances, HttpHeaders headers) {
        return predicate.getEligibleServers(serviceInstances, headers);
    }

    protected List<ServiceInstance> selectInstance(List<ServiceInstance> serviceInstances, Map<String, String> header) {
        HttpHeaders headers = new HttpHeaders();
        header.forEach((k, v) -> headers.put(k, Collections.singletonList(v)));
        return predicate.getEligibleServers(serviceInstances, headers);
    }

    @Override
    public void destroy() {
        bestAvailableWeightedSelector.shutdown();
    }
}
